Python: Preserve AG-UI tool message IDs across snapshots - #7510
Python: Preserve AG-UI tool message IDs across snapshots#7510King Star (jstar0) wants to merge 2 commits into
Conversation
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
GitHub Actions
Inactive
There was a problem hiding this comment.
🟡 Not ready to approve
Some ToolCallStartEvent emitters in _run_common.py still parent tool calls to flow.message_id while snapshots use the newly allocated segment id, which can reintroduce the stream/snapshot ID mismatch for those tool-call paths.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Pull request overview
This PR updates the Python AG-UI streaming/snapshot implementation so tool-call “assistant messages” can keep a stable identity between streamed events and the final MESSAGES_SNAPSHOT, preserving client-side ordering for interleaved text/tool-call turns.
Changes:
- Allocate and track a dedicated message ID for streamed tool-call segments and use it as
ToolCallStartEvent.parent_message_id. - Reuse the tracked tool-call segment ID when building the corresponding snapshot assistant
tool_callsmessage. - Add a regression test asserting that the snapshot tool-call message ID matches the streamed tool-call parent message ID after leading text.
File summaries
| File | Description |
|---|---|
| python/packages/ag-ui/tests/ag_ui/test_run.py | Adds regression coverage ensuring snapshot tool-call message IDs reuse the streamed tool-call parent ID. |
| python/packages/ag-ui/agent_framework_ag_ui/_run_common.py | Tracks per-tool-call-segment message IDs during streaming and uses them for ToolCallStartEvent.parent_message_id. |
| python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py | Prefers the tracked tool-call segment ID when emitting snapshot tool_calls messages. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
| def _track_tool_call_segment(flow: FlowState, tool_call_id: str) -> str: | ||
| """Record a tool call and return the message ID used by its stream events.""" | ||
| segment: dict[str, Any] | ||
| if flow.snapshot_segments and flow.snapshot_segments[-1]["kind"] == "tool_calls": | ||
| flow.snapshot_segments[-1]["call_ids"].append(tool_call_id) | ||
| segment = flow.snapshot_segments[-1] |
GitHub Actions
Inactive
| """Record a tool call in the current tool segment, opening one if needed.""" | ||
| def _new_tool_call_segment_id(flow: FlowState) -> str: | ||
| """Allocate an ID that is distinct from any streamed text segment.""" | ||
| text_message_ids = {segment.get("id") for segment in flow.snapshot_segments if segment["kind"] == "text"} |
There was a problem hiding this comment.
Could we make flow.message_id single-use once a tool-call segment claims it? With the supported tool-only preopen followed by text, or reasoning between two calls before a result, _new_tool_call_segment_id returns the same ID for distinct assistant snapshot messages. The reference client's ID-keyed MESSAGES_SNAPSHOT merge then overwrites either the text or an earlier tool call, so could we reserve the ID or allocate a fresh ID per new segment and assert snapshot-wide uniqueness?
| return generate_event_id() | ||
|
|
||
|
|
||
| def _track_tool_call_segment(flow: FlowState, tool_call_id: str) -> str: |
There was a problem hiding this comment.
Should the predictive confirm_changes emitter use this returned ID as well? _agent_run.py:2619 still discards it while _agent_run.py:2597 emits ToolCallStartEvent.parent_message_id from flow.message_id, so that approval call streams under the text message but snapshots under the tool segment. Could we mirror _emit_approval_request by tracking before emission and passing the returned ID?
| if not calls: | ||
| continue | ||
| message_id = tool_open_id or generate_event_id() | ||
| message_id = segment.get("id") or tool_open_id or generate_event_id() |
There was a problem hiding this comment.
Could we keep tool-message identity allocation local to one Module? _new_tool_call_segment_id now eagerly owns the rule in _run_common.py, but _append_segmented_snapshot_messages retains the final-state tool_open_id policy here, leaving two Modules to coordinate the same ordering invariant through the snapshot_segments Seam. That lost locality is what lets later text or reasoning invalidate the eager choice. I think one consume-once segment-ID seam makes the interface safer.
Motivation & Context
When an AG-UI run streams assistant text followed by a tool call, the stream uses the open text message ID as the tool call's parent. The final
MESSAGES_SNAPSHOTrepresents text and tool calls as separate messages, but currently assigns the tool-call message a new ID. The reference AG-UI client merges snapshots by ID, so the tool call and its result are appended after later assistant text instead of remaining in their original position.This affects any AG-UI frontend that follows the reference ID-based merge behavior and is reproducible with interleaved assistant text and tool calls.
Description & Review Guide
ToolCallStartEvent.parent_message_idand the corresponding snapshot assistant message.Related Issue
Fixes #7491
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix) — a workflow keeps the label and title prefix in sync automatically.